import gym
from pprint import pprint

env = gym.make("merge-v0")
env.configure({
  "controlled_vehicles": 2,
  "vehicles_count": 3,
  "observation": {
    "type": "MultiAgentObservation",
    "observation_config": {
      "type": "Kinematics",
    }
  },
  "action": {
    "type": "MultiAgentAction",
    "action_config": {
      "type": "DiscreteMetaAction",
    }
  }
})

pprint(env.observation_space)
pprint(env.action_space)
env.reset()
done = truncated = False
while not (done or truncated):
    action = env.action_space.sample()
    print(action)
    obs, reward, done, truncated, info = env.step(action)
    pprint(obs)
    pprint(reward)
    pprint(info)

    env.render()

# import gym
#
# env = gym.make('highway-v0')
#
# env.configure({"controlled_vehicles": 2})  # Two controlled vehicles
# env.configure({"vehicles_count": 1})  # A single other vehicle, for the sake of visualisation
# env.reset(seed=0)
#
# from matplotlib import pyplot as plt
# # %matplotlib inline
# plt.imshow(env.render(mode="rgb_array"))
# plt.title("Controlled vehicles are in green")
# plt.show()